package rabbit.flt.plugins.springmvc.plugin;

import org.springframework.http.HttpStatus;
import rabbit.flt.common.Headers;
import rabbit.flt.common.context.TraceContext;
import rabbit.flt.common.trace.MessageType;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.common.trace.io.HttpRequest;
import rabbit.flt.common.trace.io.HttpResponse;
import rabbit.flt.common.utils.StringUtils;
import rabbit.flt.plugins.common.plugin.PerformancePlugin;
import rabbit.flt.plugins.springmvc.ServletResponseWrapper;

import javax.servlet.DispatcherType;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.Enumeration;
import java.util.function.Consumer;

public class DispatcherServletPlugin extends PerformancePlugin {

    /**
     * 非200响应时缓存traceData
     */
    private static final ThreadLocal<TraceData> webTraceContext = new ThreadLocal<>();

    @Override
    public Object[] before(Object objectEnhanced, Method method, Object[] args) {
        HttpServletRequest request = (HttpServletRequest) args[0];
        if (DispatcherType.REQUEST == request.getDispatcherType()) {
            proxyResponse(args);
            webTraceContext.remove();
            TraceContext.clearContext();
            TraceContext.openTrace(method);
            String spanId = request.getHeader(Headers.SPAN_ID);
            String traceId = request.getHeader(Headers.TRACE_ID);
            if (!StringUtils.isEmpty(spanId) && !StringUtils.isEmpty(traceId)) {
                TraceContext.initRootSpanId(spanId.concat("-0"));
                TraceContext.setTraceId(traceId);
            } else {
                TraceContext.initRootSpanId("0");
            }
            super.before(objectEnhanced, method, args);
            TraceData traceData = TraceContext.getStackInfo(method).getTraceData();
            webTraceContext.set(traceData);
            setSourceApplication(request, traceData);
            traceData.setNodeName("");
            TraceContext.setWebTraceData(traceData);
        } else if (DispatcherType.ERROR == request.getDispatcherType()) {
            proxyResponse(args);
            super.before(objectEnhanced, method, args);
        }
        return args;
    }

    /**
     * 代理response
     *
     * @param args
     */
    private void proxyResponse(Object[] args) {
        HttpServletResponse response = (HttpServletResponse) args[1];
        Consumer<String> consumer = body -> {
            TraceData traceData = webTraceContext.get();
            webTraceContext.remove();
            if (null == traceData) {
                return;
            }
            setResponseInfo(args, traceData);
            if (HttpStatus.OK.value() != response.getStatus()) {
                traceData.getHttpResponse().setBody(truncate(body));
                setNodeNameWithMappedController(traceData);
            }
            traceData.updateCost();
            super.handleTraceData(traceData);
        };
        args[1] = new ServletResponseWrapper(response, consumer);
    }

    /**
     * 设置来源app字段到trace data中
     *
     * @param request
     * @param traceData
     */
    private void setSourceApplication(HttpServletRequest request, TraceData traceData) {
        String sourceApp = request.getHeader(Headers.SOURCE_APP);
        if (!StringUtils.isEmpty(sourceApp)) {
            traceData.setSourceApplication(StringUtils.toString(sourceApp));
        }
    }

    @Override
    public Object after(Object objectEnhanced, Method method, Object[] args, Object result) {
        HttpServletRequest request = (HttpServletRequest) args[0];
        if (DispatcherType.ERROR == request.getDispatcherType()) {
            return result;
        }
        return super.after(objectEnhanced, method, args, result);
    }

    @Override
    public void onException(Object objectEnhanced, Method method, Object[] args, Throwable t) {
        HttpServletRequest request = (HttpServletRequest) args[0];
        if (DispatcherType.ERROR == request.getDispatcherType()) {
            return;
        }
        super.onException(objectEnhanced, method, args, t);
    }

    @Override
    protected void fillTraceData(TraceData traceData, Object objectEnhanced, Method method, Object[] args, Object result) {
        HttpServletRequest request = (HttpServletRequest) args[0];
        if (StringUtils.isEmpty(traceData.getNodeName())) {
            traceData.setNodeName(request.getRequestURI());
        } else {
            // 如果下游填充了path，则补全context path
            if (!"/".equals(request.getContextPath().trim())) {
                traceData.setNodeName(request.getContextPath().concat(traceData.getNodeName()));
            }
        }
        traceData.setNodeDesc(request.getMethod().concat(": ").concat(traceData.getNodeName()));
        traceData.setMessageType(MessageType.REST.name());
        setRequestInfo(args, traceData);
        setResponseInfo(args, traceData);
    }

    /**
     * 设置请求参数
     *
     * @param args
     * @param traceData
     */
    private void setRequestInfo(Object[] args, TraceData traceData) {
        HttpServletRequest request = (HttpServletRequest) args[0];
        HttpRequest requestInfo = new HttpRequest();
        requestInfo.setContextPath(request.getContextPath());
        requestInfo.setMethod(request.getMethod());
        requestInfo.setRequestUri(request.getRequestURI());
        Enumeration<String> names = request.getHeaderNames();
        while (names.hasMoreElements()) {
            String name = names.nextElement();
            requestInfo.addHeader(name, truncate(request.getHeader(name)));
        }
        names = request.getParameterNames();
        while (names.hasMoreElements()) {
            String name = names.nextElement();
            requestInfo.addParameter(name, truncate(request.getParameter(name)));
        }
        traceData.setHttpRequest(requestInfo);
    }

    /**
     * 设置响应
     *
     * @param args
     * @param traceData
     */
    private void setResponseInfo(Object[] args, TraceData traceData) {
        HttpServletResponse response = (HttpServletResponse) args[1];
        if (null != traceData.getHttpResponse()) {
            return;
        }
        HttpResponse responseInfo = new HttpResponse();
        responseInfo.setStatusCode(response.getStatus());
        for (String name : response.getHeaderNames()) {
            responseInfo.addHeader(name, truncate(response.getHeader(name)));
        }
        traceData.setHttpResponse(responseInfo);
    }

    @Override
    public void doFinal(Object objectEnhanced, Method method, Object[] args, Object result) {
        HttpServletRequest request = (HttpServletRequest) args[0];
        if (DispatcherType.ERROR == request.getDispatcherType()) {
            return;
        }
        super.doFinal(objectEnhanced, method, args, result);
    }

    @Override
    protected void handleTraceData(TraceData traceData) {
        // do nothing, 放在填充数据阶段发送trace
    }

}
